package pro.softzhang.algo.lc1100;

import pro.softzhang.algo.flag.theme.TDesign;

import java.util.*;

/**
 * 设计题
 */
@TDesign
public class LC1172_DinnerPlateStacks {
    public static void main(String[] args) {

    }

    static
    class DinnerPlates {
        private int capacity;
        private List<Stack<Integer>> stacks = new ArrayList<>();
        private TreeSet<Integer> notFull = new TreeSet<>();

        public DinnerPlates(int capacity) {
            this.capacity = capacity;
        }

        public void push(int val) {
            if (notFull.isEmpty()) {
                Stack newStack = new Stack<>();
                newStack.push(val);
                stacks.add(newStack);
                if (capacity > 1) {
                    notFull.add(stacks.size() - 1);
                }
                return;
            }

            int index = notFull.first();
            stacks.get(index).push(val);
            if (stacks.get(index).size() == capacity) {
                notFull.pollFirst();
            }
        }

        public int pop() {
            return popAtStack(stacks.size() - 1);
        }

        public int popAtStack(int index) {
            if (index < 0 || index >= stacks.size() || stacks.get(index).isEmpty()) {
                return -1;
            }
            int val = stacks.get(index).pop();
            if (index == stacks.size() - 1 && stacks.get(stacks.size() - 1).isEmpty()) {
                while (!stacks.isEmpty() && stacks.get(stacks.size() - 1).isEmpty()) {
                    notFull.remove(stacks.size() - 1);
                    stacks.remove(stacks.size() - 1);
                }
            } else {
                notFull.add(index);
            }
            return val;
        }
    }
}
